import torch
import pandas as pd
import numpy as np
import torch.nn as nn
from sklearn.preprocessing import StandardScaler
import joblib

# 加载测试数据
test_df = pd.read_csv('processed_clt_6_data.csv')

# 定义模型结构
class ResNetLikeCNN_DNN(nn.Module):
    def __init__(self, input_size):
        super(ResNetLikeCNN_DNN, self).__init__()
        # 第一层卷积
        self.conv1 = nn.Conv1d(1, 16, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm1d(16)
        self.relu1 = nn.LeakyReLU(0.01)

        # ResNet 块 1
        self.resblock1_conv1 = nn.Conv1d(16, 16, kernel_size=3, padding=1)
        self.resblock1_bn1 = nn.BatchNorm1d(16)
        self.resblock1_conv2 = nn.Conv1d(16, 16, kernel_size=3, padding=1)
        self.resblock1_bn2 = nn.BatchNorm1d(16)
        self.resblock1_relu = nn.LeakyReLU(0.01)

        # ResNet 块 2
        self.resblock2_conv1 = nn.Conv1d(16, 32, kernel_size=3, padding=1)
        self.resblock2_bn1 = nn.BatchNorm1d(32)
        self.resblock2_conv2 = nn.Conv1d(32, 32, kernel_size=3, padding=1)
        self.resblock2_bn2 = nn.BatchNorm1d(32)
        self.resblock2_relu = nn.LeakyReLU(0.01)

        # 添加投影操作
        self.projection2 = nn.Conv1d(16, 32, kernel_size=1)

        # ResNet 块 3
        self.resblock3_conv1 = nn.Conv1d(32, 64, kernel_size=3, padding=1, stride=2)
        self.resblock3_bn1 = nn.BatchNorm1d(64)
        self.resblock3_conv2 = nn.Conv1d(64, 64, kernel_size=3, padding=1)
        self.resblock3_bn2 = nn.BatchNorm1d(64)
        self.resblock3_relu = nn.LeakyReLU(0.01)

        # 添加投影操作
        self.projection3 = nn.Conv1d(32, 64, kernel_size=1, stride=2)

        self.fc1 = nn.Linear(64 * input_size // 2, 256)  # 注意这里的input_size需要调整
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 64)
        self.fc4 = nn.Linear(64, 1)
        self.relu = nn.LeakyReLU(0.01)
        self.dropout = nn.Dropout(p=0.4)

    def forward(self, x):
        # 第一层
        x1 = self.relu1(self.bn1(self.conv1(x)))
        # print("x1 shape:", x1.shape)

        # ResNet 块 1
        res1 = self.resblock1_conv1(x1)
        res1 = self.resblock1_bn1(res1)
        res1 = self.resblock1_relu(res1)
        res1 = self.resblock1_conv2(res1)
        res1 = self.resblock1_bn2(res1)
        x2 = self.resblock1_relu(x1 + res1)
        # print("x2 shape:", x2.shape)

        # ResNet 块 2
        res2 = self.resblock2_conv1(x2)
        res2 = self.resblock2_bn1(res2)
        res2 = self.resblock2_relu(res2)
        res2 = self.resblock2_conv2(res2)
        res2 = self.resblock2_bn2(res2)
        # print("res2 shape:", res2.shape)

        # 投影操作
        x2_projected = self.projection2(x2)
        # print("x2_projected shape:", x2_projected.shape)

        x3 = self.resblock2_relu(x2_projected + res2)
        # print("x3 shape:", x3.shape)

        # 投影操作 for ResNet 块 3
        x3_projected = self.projection3(x3)
        # print("x3_projected shape:", x3_projected.shape)

        # ResNet 块 3
        res3 = self.resblock3_conv1(x3)
        res3 = self.resblock3_bn1(res3)
        res3 = self.resblock3_relu(res3)
        res3 = self.resblock3_conv2(res3)
        res3 = self.resblock3_bn2(res3)
        x4 = self.resblock3_relu(x3_projected + res3)
        # print("x4 shape:", x4.shape)

        x4 = x4.view(x4.size(0), -1)
        x5 = self.relu(self.fc1(x4))
        x5 = self.dropout(x5)
        x6 = self.relu(self.fc2(x5))
        x7 = self.relu(self.fc3(x6))
        x8 = self.fc4(x7)
        return x8
# 设置设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 定义要使用的特征列
features = [
    'fy_cth_scaled', 'fy_ctt_scaled', 'fy_ctp_scaled', 'fy_olr_scaled',
    'fy_lat_sin', 'fy_lat_cos', 'fy_lon_sin', 'fy_lon_cos',
    'band1_scaled', 'band2_scaled', 'band3_scaled', 'band4_scaled',
    'band5_scaled', 'band6_scaled', 'band7_brightness_temperature', 'band7_scaled',
    'band8_brightness_temperature', 'band8_scaled', 'band9_brightness_temperature', 'band9_scaled',
    'band10_brightness_temperature', 'band10_scaled', 'band11_brightness_temperature', 'band11_scaled',
    'band12_brightness_temperature', 'band12_scaled', 'band13_brightness_temperature', 'band13_scaled',
    'band14_brightness_temperature', 'band14_scaled', 'band1_gray_value_scaled', 'band2_gray_value_scaled',
    'band3_gray_value_scaled', 'band4_gray_value_scaled', 'band5_gray_value_scaled', 'band6_gray_value_scaled',
    'band7_gray_value_scaled', 'band8_gray_value_scaled', 'band9_gray_value_scaled', 'band10_gray_value_scaled',
    'band11_gray_value_scaled', 'band12_gray_value_scaled', 'band13_gray_value_scaled', 'band14_gray_value_scaled'
]

# 提取测试数据的特征
X_test = test_df[features].values

# 将特征标准化
scaler_x = StandardScaler()
X_test_scaled = scaler_x.fit_transform(X_test)

# 转换为张量
X_test_tensor = torch.tensor(X_test_scaled, dtype=torch.float32)

# 初始化模型并加载权重
model = ResNetLikeCNN_DNN(X_test.shape[1]).to(device)
model.load_state_dict(torch.load('trained_model_ocean6.pth'))
model.eval()  # 设置为评估模式

# 设置批量大小
batch_size = 64  # 根据显存大小调整

# 分批进行预测
y_test_pred = []
with torch.no_grad():
    for i in range(0, X_test_tensor.size(0), batch_size):
        batch_X = X_test_tensor[i:i + batch_size].unsqueeze(1).to(device)
        batch_pred = model(batch_X).cpu().numpy()
        y_test_pred.append(batch_pred)

# 合并所有批次的预测结果
y_test_pred = np.concatenate(y_test_pred).flatten()

# 加载训练阶段保存的 scaler_y 对象
scaler_y = joblib.load('scaler6_y.pkl')

# 反标准化目标值
y_test_pred_true = scaler_y.inverse_transform(y_test_pred.reshape(-1, 1)).flatten()

# 创建包含预测值的 DataFrame
# test_result_df = pd.DataFrame({'Predicted': y_test_pred_true})
# 创建包含预测值和原始数据的 DataFrame
test_result_df = pd.DataFrame({
    'fy_cbh': y_test_pred_true,
    'fy_cth': test_df['fy_cth'].values,
    'fy_clt': test_df['fy_clt'].values,
    'fy_lat': test_df['fy_lat'].values,
    'fy_lon': test_df['fy_lon'].values,
    'fy_cot': test_df['fy_cth'].values - y_test_pred_true
})
# jieguo
# 保存预测结果至 CSV 文件
test_result_df.to_csv('prediction_ocean6_2020010104_results.csv', index=False)
